#!/home/handuser/miniconda3/envs/yjh/bin/python3
import numpy as np
import os, sys
from sensor_msgs.msg import Image
import rospy
import imgviz

sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')
HERE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(HERE_DIR, '../../instance_segmentation/'))
from mmdet.apis import inference_detector, init_detector
sys.path.append('/opt/ros/kinetic/lib/python2.7/dist-packages')

sys.path.append(os.path.join(HERE_DIR, '../../utils/'))
from image_converter import convert_Image_to_nparray, convert_nparray_to_Image

# cat       coco_id     nocs_id
# mug       41          5
# bowl      45          1
# bottle    39          0
coco_list = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
            'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
            'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
            'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
            'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
            'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
            'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
            'bottle', 'wine glass', 'mug', 'fork', 'knife', 'spoon', 'bowl',
            'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
            'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
            'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
            'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
            'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
            'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
id_map = {
    41: 5,
    45: 1,
    39: 0
}

def get_bbox_from_mask(mask):
    # mask: bool array, shape: (H, W)
    # return: (y1, x1, y2, x2)
    mask_ids = np.argwhere(mask == True)
    left_top = mask_ids.min(axis=0)
    right_bottom = mask_ids.max(axis=0)
    return np.array([left_top[0], left_top[1], right_bottom[0], right_bottom[1]])

class InstSeg:
    def __init__(self, device) -> None:
        here = os.path.dirname(os.path.abspath(__file__))
        cfg = os.path.join(here, '../configs/solo/decoupled_solo_light_r50_fpn_3x_coco.py')
        ckpt = os.path.join(here, '../assets/decoupled_solo_weights.pth')
        self.device = device
        self.score_thr = 0.5
        self.model = init_detector(cfg, ckpt, device)

    def serve_as_regular_publisher(self, rate_hz=2):
        rospy.init_node('instance_segmentaion', anonymous=True)
        self.puber = rospy.Publisher('/dexgrasp/mask', Image, queue_size=1)
        rate = rospy.Rate(rate_hz)
        while not rospy.is_shutdown():
            img_raw = rospy.wait_for_message("/camera/color/image_raw", Image)
            _, img_masked_msg = self.run(img_raw)
            self.puber.publish(img_masked_msg)
            rate.sleep()

    def run(self, image_raw:Image):
        img = convert_Image_to_nparray(image_raw)
        result = inference_detector(self.model, img)
        result_compressed = {'class_names':[], 'class_ids':[], 'scores':[], 'masks': [], 'bboxes':[]}
        for id in id_map.keys():  # show mug, bottle and bowl
        # for id in range(79, -1, -1): # show all categories
            if result[0][id].shape[0] > 0:
                # >>>>>> choice 1 : choose best detection results
                # score = result[0][id][:, 4].max()
                # if score > self.score_thr:
                #     cat_id = id                
                #     result_compressed['class_names'].append(coco_list[cat_id])
                #     result_compressed['class_ids'].append(cat_id)  # raw coco
                #     # result_compressed['class_ids'].append(id_map[cat_id])  # coco to nocs
                #     best_idx = result[0][cat_id][:, 4].argmax()
                #     result_compressed['scores'].append(score)
                #     mask = np.array(result[1][cat_id])[best_idx]
                #     result_compressed['masks'].append(mask)
                #     result_compressed['bboxes'].append(get_bbox_from_mask(mask))      
                # >>>>>> choice 2 : choose all results whoose scores greater than threshold          
                for i in range(0, result[0][id].shape[0]):
                    score = result[0][id][i, 4]
                    if score > self.score_thr:
                        cat_id = id
                        result_compressed['class_names'].append(coco_list[cat_id])
                        # result_compressed['class_ids'].append(cat_id) # raw coco
                        result_compressed['class_ids'].append(id_map[cat_id]) # coco to nocs
                        result_compressed['scores'].append(score)
                        mask = np.array(result[1][cat_id])[i]
                        result_compressed['masks'].append(mask)
                        result_compressed['bboxes'].append(get_bbox_from_mask(mask))
        if len(result_compressed['class_ids']) == 0:
            return None, image_raw
        for k, v in result_compressed.items():
            result_compressed[k] = np.array(v)

        labels = result_compressed['class_ids']

        captions = result_compressed['class_names']
        masks = result_compressed['masks']
        img_masked = imgviz.instances2rgb(img, masks=masks, labels=labels, captions=captions)
        img_masked_msg = convert_nparray_to_Image(img_masked)
        
        return result_compressed, img_masked_msg